﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Net.Sockets;
using System.Threading;
using System.IO;

namespace TcpProxy
{
    public class TcpProxy
    {
        const int DefaultForceTimeoutMs = 2000;

        const int DelaySleepTime = 10;

        string _targetHost = "";
        int _targetPort = 0;
        int _sourcePort;

        bool _shouldExit = false;
        bool _isActive = false;
        bool _reListen = false;

        BlockingStreamReader.Endianness _endianness;

        Thread _proxyThread;

        TcpClient _serverTcpClient, _clientTcpClient;

        ProtocolDecoder _serverProtocolDecoder, _clientProtocolDecoder;

        System.IO.Pipes.AnonymousPipeServerStream _serverPipe, _clientPipe;

        public string ServerPacketDumpFile { get; set; }
        public string ClientPacketDumpFile { get; set; }

        public event EventHandler ProxySessionStarted;
        public event EventHandler ProxySessionFinished;

        public ProtocolDecoder ServerProtocolDecoder { get { return _serverProtocolDecoder; } }
        public ProtocolDecoder ClientProtocolDecoder { get { return _clientProtocolDecoder; } }

        public void CloseTcpClients()
        {
            _reListen = true;
        }

        public void SetProtocolDecoders(ProtocolDecoder serverProtocolDecoder, ProtocolDecoder clientProtocolDecoder)
        {
            closePipes();

            //create fresh pipes
            _serverPipe = new System.IO.Pipes.AnonymousPipeServerStream();
            _clientPipe = new System.IO.Pipes.AnonymousPipeServerStream();
            
            serverProtocolDecoder.PacketStream = new BlockingStreamReader(
                new System.IO.Pipes.AnonymousPipeClientStream(System.IO.Pipes.PipeDirection.In, _serverPipe.ClientSafePipeHandle),
                _endianness);

            clientProtocolDecoder.PacketStream = new BlockingStreamReader(
                new System.IO.Pipes.AnonymousPipeClientStream(System.IO.Pipes.PipeDirection.In, _clientPipe.ClientSafePipeHandle),
                _endianness);

            serverProtocolDecoder.TcpProxy = this;
            clientProtocolDecoder.TcpProxy = this;

            _serverProtocolDecoder = serverProtocolDecoder;
            _clientProtocolDecoder = clientProtocolDecoder;
        }

        public void ForceExit()
        {
            _shouldExit = true;
        }

        public bool IsActive { get { return _isActive; } }

        public TcpProxy(string targetHost, int targetPort, int sourcePort, BlockingStreamReader.Endianness endianness)
        {
            _targetHost = targetHost;
            _targetPort = targetPort;
            _sourcePort = sourcePort;

            _endianness = endianness;

        }

        public void Start()
        {
            if (_isActive) throw new Exception("Proxy is already active!");

            _proxyThread = new Thread(new ThreadStart(proxyThreadMain));
            _proxyThread.Name = "Tcp Proxy";
            _proxyThread.Start();

        }

        public void Stop()
        {
            Stop(DefaultForceTimeoutMs);
        }

        public void Stop(int forceTimeoutMs)
        {
            if (!_isActive) return;
            if (forceTimeoutMs <= 0) throw new Exception("Invalid timeout!");

            _shouldExit = true;
            DateTime startTime = DateTime.Now;
            while (_isActive)
            {
                System.Threading.Thread.Sleep(100);
                if (DateTime.Now.Subtract(startTime).TotalMilliseconds > forceTimeoutMs)
                {
                    _proxyThread.Abort();
                    _isActive = false;
                    return;
                }
            }

            closePipes();
        }

        private void proxyThreadMain()
        {
            try
            {
                _shouldExit = false;
                _isActive = true;

                Console.WriteLine("TcpProxy starting");

                while (!_shouldExit)
                {
                    Console.WriteLine(string.Format("TcpProxy listening on localhost:{0}", _sourcePort));
                    TcpListener listener = new TcpListener(System.Net.IPAddress.Any, _sourcePort);
                    listener.Start();

                    while (!listener.Pending())
                    {
                        if (_shouldExit)
                        {
                            listener.Stop();
                            return;
                        }
                        System.Threading.Thread.Sleep(DelaySleepTime);
                    }

                    _clientTcpClient = listener.AcceptTcpClient();
                    listener.Stop();
                    _clientTcpClient.NoDelay = true;

                    Console.WriteLine(string.Format("TcpProxy accepted connection from {0}", _clientTcpClient.Client.RemoteEndPoint.ToString()));

                    _serverTcpClient = new TcpClient();
                    _serverTcpClient.Connect(_targetHost, _targetPort);
                    _serverTcpClient.NoDelay = true;

                    Console.WriteLine(string.Format("TcpProxy created connection to {0}:{1}", _targetHost, _targetPort));

                    if (ProxySessionStarted != null)
                        ProxySessionStarted(this, new EventArgs());

                    _reListen = false;

                    while (
                        (_serverTcpClient.Client != null) && (_clientTcpClient.Client != null) &&
                        (_serverTcpClient.Client.Connected) && (_clientTcpClient.Client.Connected) && (!_shouldExit) &&
                        (!_reListen))
                    {
                        if ((_serverTcpClient.Client != null) && (_serverTcpClient.Available > 0))
                        {
                            proxyData(_serverTcpClient, _clientTcpClient, _serverPipe);
                        }

                        if ((_clientTcpClient.Client != null) && (_clientTcpClient.Available > 0))
                        {
                            proxyData(_clientTcpClient, _serverTcpClient, _clientPipe);
                        }

                        System.Threading.Thread.Sleep(DelaySleepTime);
                    }

                    Console.WriteLine("TcpProxy closing connections");

                    try
                    {
                        _clientTcpClient.Close();
                    }
                    catch { }

                    try
                    {
                        _serverTcpClient.Close();
                    }
                    catch { }

                    if (ProxySessionFinished != null)
                        ProxySessionFinished(this, new EventArgs());
                }

                Console.WriteLine("TcpProxy Stopped");
            }
            finally
            {
                _isActive = false;
            }
        }

        private void proxyData(TcpClient fromClient, TcpClient toClient, System.IO.Pipes.AnonymousPipeServerStream pipe)
        {

            byte[] buff = new byte[fromClient.Available];

            fromClient.Client.Receive(buff);

            toClient.Client.Send(buff);

            string path = (fromClient == _clientTcpClient) ? ClientPacketDumpFile : ServerPacketDumpFile;

            if ((path != null) && (path != ""))
            {
                FileStream stream = System.IO.File.Open(path, System.IO.FileMode.OpenOrCreate);
                stream.Position = stream.Length;

                stream.Write(buff, 0, buff.Length);
                stream.Close();
            }

            if (pipe != null)
            {
                pipe.Write(buff, 0, buff.Length);
                pipe.Flush();
            }
        }

        private void closePipes()
        {
            try
            {
                if (_serverPipe != null) _serverPipe.Close();
            }
            catch { }

            try
            {
                if (_clientPipe != null) _clientPipe.Close();
            }
            catch { }
        }
    }
}
